import torch
import csv
import os
import math

import smtplib
from email.mime.text import MIMEText
from email.mime.image import MIMEImage
from email.mime.multipart import MIMEMultipart
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

from torchvision import transforms, datasets

device = torch.device('cuda' if torch.cuda.is_available else 'cpu')

train_dataset = datasets.MNIST(root='../../data', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.MNIST(root='../../data', train=False, transform=transforms.ToTensor(), download=True)


def getdata():
    train_x = []
    train_y = []
    for data in train_dataset:
        train_x.append(data[0].tolist())
        train_y.append(data[1])
    train_x = torch.tensor(train_x).squeeze().view(len(train_x), -1)
    train_y = torch.tensor(train_y)
    print(train_x.shape, train_y.shape)

    test_x = []
    test_y = []
    for data in test_dataset:
        test_x.append(data[0].tolist())
        test_y.append(data[1])
    test_x = torch.tensor(test_x).squeeze().view(len(test_x), -1)
    test_y = torch.tensor(test_y)
    print(test_x.shape, test_y.shape)

    return train_x.to(device), train_y.to(device), test_x.to(device), test_y.to(device)


def knn(_x, x, y, k):
    temp = _x - x
    temp = torch.pow(temp, 2)
    temp = temp.sum(dim=1)
    temp = torch.sqrt(temp)
    argsort = temp.argsort()  # [3 2 1 0]
    result = y[argsort][:k]  # [1 1 0]
    result = result.int()
    return torch.bincount(result).argmax()  # 取众数


ks = []


def start():
    if not os.path.exists('../csv'):
        os.mkdir('../csv')
    log_path = '../csv/k.csv'
    file = open(log_path, 'a+', encoding='utf-8', newline='')
    csv_writer = csv.writer(file)
    csv_writer.writerow(['k', 'accuracy'])

    x, y, test_x, test_y = getdata()

    for k in range(1, 301):
        correct = 0
        for i in range(len(test_x)):
            pred = knn(test_x[i], x, y, k=k)
            if pred == test_y[i]:
                correct += 1
            print(f'k={k},step:{i + 1}/{len(test_x)},pred={pred},y={test_y[i]},correct={correct / (i + 1)}')
        print('k={},correct={}'.format(k, correct / len(test_x)))
        csv_writer.writerow([k, correct / len(test_x)])
        ks.append((correct / len(test_x), k))
    file.close()


class MyEmail():
    def __init__(self, subject):
        self.host = 'smtp.qq.com'
        self.port = 465
        self.user = 'xueguopeng@foxmail.com'
        self.password = 'islclqzazvnidbdc'
        self.subject = subject

    def send_good_email(self):
        if len(ks) != 0:
            msg = '<p>准确率最高的k值: ' + str(sorted(ks, reverse=True)[:3]) \
                  + '</p><img src="cid:k_img">'
        else:
            msg = 'train_accuracy_s or test_accuracy_s is None'

        subject = self.subject + '-脚本运行结束通知'

        email_client = smtplib.SMTP_SSL(host=self.host)
        # 设置发件人邮箱的域名和端口
        email_client.connect(host=self.host, port=self.port)

        # 登陆邮件，权限验证，password为邮箱密码
        result = email_client.login(user=self.user, password=self.password)
        print("登录结果", result)

        # 发送邮件，from_addr：发送人，to_addrs：收件人 ，msg：发送的文本
        message = MIMEMultipart()
        message['From'] = self.user
        message['To'] = self.user
        message['Subject'] = subject

        # 发送图片，增加图片标签
        if os.path.exists('../img/accuracy-k.png'):
            with open('../img/accuracy-k.png', 'rb') as f:
                image_data = f.read()
            image = MIMEImage(image_data)
            image.add_header('Content-ID', '<k_img>')
            message.attach(image)

        message.attach(MIMEText(msg, 'html'))
        email_client.sendmail(from_addr=self.user, to_addrs=self.user, msg=message.as_string())
        # 关闭邮件客户端
        email_client.close()

    def send_bad_email(self, msg):

        subject = self.subject + '-脚本运行失败通知'

        email_client = smtplib.SMTP_SSL(host=self.host)
        # 设置发件人邮箱的域名和端口
        email_client.connect(host=self.host, port=self.port)

        # 登陆邮件，权限验证，password为邮箱密码
        result = email_client.login(user=self.user, password=self.password)
        print("登录结果", result)

        # 发送邮件，from_addr：发送人，to_addrs：收件人 ，msg：发送的文本
        message = MIMEMultipart()
        message['From'] = self.user
        message['To'] = self.user
        message['Subject'] = subject
        message.attach(MIMEText(msg))
        email_client.sendmail(from_addr=self.user, to_addrs=self.user, msg=message.as_string())
        # 关闭邮件客户端
        email_client.close()


def analyze():
    if os.path.exists('../csv/k.csv'):
        k_res = pd.read_csv('../csv/k.csv')
        k = k_res['k']
        acc = k_res['accuracy']

        plt.rcParams['figure.figsize'] = (12.8, 7.2)  # 1280 x 720 像素全局设置
        new_ticks = np.arange(1, len(k), math.ceil(len(k) / 30))

        k_acc_curve, = plt.plot(k, acc)
        plt.xlabel('k')
        plt.ylabel('accuracy')
        plt.title('accuracy-k')
        plt.xticks(new_ticks)
        if not os.path.exists('../img'):
            os.mkdir('../img')
        plt.savefig('../img/accuracy-k.png')
        plt.close()


if __name__ == '__main__':
    email = MyEmail('基于knn算法的手写数字识别')
    try:
        print('开始测试啦')
        start()
        print('开始分析啦')
        analyze()
    except:
        print('发送失败邮件啦')
        email.send_bad_email('有bug')
    else:
        print('发送成功邮件啦')
        email.send_good_email()
